/*
* crypto sha and aes driver test module
* Licensed under GPLv2 or later.
*/

#include <crypto/hash.h>
#include <linux/err.h>
#include <linux/init.h>
#include <linux/gfp.h>
#include <linux/module.h>
#include <linux/scatterlist.h>
#include <linux/string.h>
#include <linux/moduleparam.h>
#include <linux/jiffies.h>
#include <linux/timex.h>
#include <linux/interrupt.h>

#define test_dbg(fmt,args...)	printk(fmt, ##args)

static struct completion com;

static void encryption_test_complete(struct crypto_async_request *r, 
	int err)
{
	struct ablkcipher_request *req;
	int i;
	struct scatterlist *sg;
	unsigned char *p;
	int n = 1;
	req = ablkcipher_request_cast(r);
	sg = req->dst;

	test_dbg("ciphertext:\n");
	while(1)
	{
		test_dbg("sg%d: len = %d\n", n, sg_dma_len(sg));
		p = sg_virt(sg);
		for(i = 0; i < sg_dma_len(sg); i++)
		{
			test_dbg("%x ", *(p + i));
		}
		test_dbg("\n");

		if(sg_is_last(sg))
			break;

		sg = sg_next(sg);
		n++;
	}

	test_dbg("encryption_test_complete err = %d\n", err);
	complete(&com);
}


static void decryption_test_complete(struct crypto_async_request *r,
	int err)
{
	struct ablkcipher_request *req;
	int i;
	struct scatterlist *sg;
	char *p;
	int n = 1;
	req = ablkcipher_request_cast(r);
	sg = req->dst;

	test_dbg("plaintext:\n");
	while(1)
	{
		test_dbg("sg%d: len = %d\n", n, sg_dma_len(sg));
		p = sg_virt(sg);
		for(i = 0; i < sg_dma_len(sg); i++)
		{
			test_dbg("%c", *(p + i));
		}
		test_dbg("\n");

		if(sg_is_last(sg))
			break;

		sg = sg_next(sg);
		n++;
	}

	test_dbg("decryption_test_complete err = %d\n", err);
	complete(&com);
}

static char plaintext[128] __attribute__((unused, aligned(8))) = 
{"12390jfdsajlj4320123456789abcdef6789abcdcbafjdlatfdsafdsafjfzcv38984fdsf23325tdsfddhfdbnvbmkjlag"};

static char plaintext1[128] __attribute__((unused, aligned(8))) = 
	{"12390jfdsajlj432143232fsfdsfdsaa"};
static char plaintext2[128] __attribute__((unused, aligned(8))) = 
	{"0123456789abcdeffedcba9876543210"};
static char plaintext3[128] __attribute__((unused, aligned(8))) = 
	{"6789abcdcbafjdlaskjfdklsaio13921"};
static char plaintext4[128] __attribute__((unused, aligned(8))) = 
	{"tfdsafdsafjfzcv3123fasfliuouprez"};	
static char plaintext5[128] __attribute__((unused, aligned(8))) = 
	{"8984fdsf23325tdsfddhfdbnvbmkjlag"};

static char ciphertext[128] __attribute__((unused, aligned(8)));
static char ciphertext1[128] __attribute__((unused, aligned(8)));
static char ciphertext2[128] __attribute__((unused, aligned(8)));
static char ciphertext3[128] __attribute__((unused, aligned(8)));
static char ciphertext4[128] __attribute__((unused, aligned(8)));
static char ciphertext5[128] __attribute__((unused, aligned(8)));

static int crypto_aes_test(void)
{
	struct scatterlist sg_src[8], sg_dst[8];
	char result[128];
	struct crypto_ablkcipher *tfm;
	struct ablkcipher_request *req;
	char *iv = "0123456789abcdef";
	char *key = "---9876543210---";
	int ret;
	int len;
	
	test_dbg("crypto_test_init\n");

	tfm = crypto_alloc_ablkcipher("ctr(aes)", 0, 0);
	if (IS_ERR(tfm)) {
		printk(KERN_ERR "alg: skcipher: Failed to load transform for"
		       "%ld\n", PTR_ERR(tfm));
		return PTR_ERR(tfm);
	}

	req = ablkcipher_request_alloc(tfm, GFP_KERNEL);
	if (!req) {
		printk(KERN_ERR "alg: skcipher: Failed to allocate request\n");
		return -1;
	}

	ablkcipher_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
					encryption_test_complete, &result);

	crypto_ablkcipher_setkey(tfm, key, 16);

	test_dbg("tfm = %p	req = %p\n", tfm, req);


	//test aes encryption with scatterlist table only having one scatterlist. 
	len = 96;
	sg_init_one(sg_src, plaintext, 96);
	sg_init_one(sg_dst, ciphertext, 96);

	ablkcipher_request_set_crypt(req, sg_src, sg_dst, len, iv);
	ret = crypto_ablkcipher_encrypt(req);

	init_completion(&com);
	wait_for_completion(&com);

	//test aes encryption with scatterlist table having more than one scatterlist. 
	len = 16 + 16 + 16 + 16 + 32;

	sg_init_table(sg_src, 5);
	sg_set_buf(sg_src + 0, plaintext1, 16);
	sg_set_buf(sg_src + 1, plaintext2, 16);
	sg_set_buf(sg_src + 2, plaintext3, 16);
	sg_set_buf(sg_src + 3, plaintext4, 16);
	sg_set_buf(sg_src + 4, plaintext5, 32);

	sg_init_table(sg_dst, 5);
	sg_set_buf(sg_dst + 0, ciphertext1, 16);
	sg_set_buf(sg_dst + 1, ciphertext2, 16);
	sg_set_buf(sg_dst + 2, ciphertext3, 16);
	sg_set_buf(sg_dst + 3, ciphertext4, 16);
	sg_set_buf(sg_dst + 4, ciphertext5, 32);

	ablkcipher_request_set_crypt(req, sg_src, sg_dst, len, iv);
	ret = crypto_ablkcipher_encrypt(req);

	init_completion(&com);
	wait_for_completion(&com);

	test_dbg("aes encryption test done.\n");


	memset(plaintext, 0, 128);
	memset(plaintext1, 0, 128);
	memset(plaintext2, 0, 128);
	memset(plaintext3, 0, 128);	
	memset(plaintext4, 0, 128);
	memset(plaintext5, 0, 128);

	ablkcipher_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
					decryption_test_complete, &result);

	crypto_ablkcipher_setkey(tfm, key, 16);

	//test aes decryption with scatterlist table only having one scatterlist. 
	len = 96;
	sg_init_one(sg_src, plaintext, 96);
	sg_init_one(sg_dst, ciphertext, 96);

	ablkcipher_request_set_crypt(req, sg_dst, sg_src, len, iv);
	ret = crypto_ablkcipher_encrypt(req);

	init_completion(&com);
	wait_for_completion(&com);


	//test aes decryption with scatterlist table having more than one scatterlist. 
	len = 16 + 8 + 8 + 16 + 32;

	sg_init_table(sg_src, 5);
	sg_set_buf(sg_src + 0, plaintext1, 16);
	sg_set_buf(sg_src + 1, plaintext2, 16);
	sg_set_buf(sg_src + 2, plaintext3, 16);
	sg_set_buf(sg_src + 3, plaintext4, 16);
	sg_set_buf(sg_src + 4, plaintext5, 32);

	sg_init_table(sg_dst, 5);
	sg_set_buf(sg_dst + 0, ciphertext1, 16);
	sg_set_buf(sg_dst + 1, ciphertext2, 16);
	sg_set_buf(sg_dst + 2, ciphertext3, 16);
	sg_set_buf(sg_dst + 3, ciphertext4, 16);
	sg_set_buf(sg_dst + 4, ciphertext5, 32);

	ablkcipher_request_set_crypt(req, sg_dst, sg_src, len, iv);
	ret = crypto_ablkcipher_encrypt(req);

	init_completion(&com);
	wait_for_completion(&com);

	test_dbg("aes decryption test done.\n");

	ablkcipher_request_free(req);
	crypto_free_ablkcipher(tfm);
	return 0;
}

static void sha_test_complete(struct crypto_async_request *r, int err)
{
	struct ahash_request *req = ahash_request_cast(r);
	int i;

	test_dbg("sha hash:\n");
	for(i = 0; i < 32; i++)
	{
		test_dbg("%2x ", *(req->result + i));
	}
	
	test_dbg("\nsha_test_complete err = %d\n", err);
	complete(&com);
}

static int crypto_sha_test(void)
{
	struct crypto_ahash *tfm;
	struct ahash_request *req;
	struct scatterlist sg_src[8];
	int len;
	u8 result[32];

	tfm = crypto_alloc_ahash("sha256", 0, 0);
	if (IS_ERR(tfm)) {
		printk(KERN_ERR "failed to load transform \n");
		return -1;
	}

	req = ahash_request_alloc(tfm, GFP_KERNEL);
	if (!req) {
		printk(KERN_ERR "alg: hash: Failed to allocate request\n");
		return -1;
	}
	ahash_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
				   sha_test_complete, NULL);

	memset(result, 0, 32);

	//test sha algorithm with scatterlist table having only one scatterlist. 
	len = 96;
	sg_init_one(sg_src, plaintext, 96);
	
	ahash_request_set_crypt(req, sg_src, result, len);
	crypto_ahash_digest(req);

	init_completion(&com);
	wait_for_completion(&com);


	memset(result, 0, 32);

	//test sha algorithm with scatterlist table having more than one scatterlist. 
	len = 16 + 8 + 8 + 16 + 8;
	sg_init_table(sg_src, 5);
	sg_set_buf(sg_src + 0, plaintext1, 16);
	sg_set_buf(sg_src + 1, plaintext2, 8);
	sg_set_buf(sg_src + 2, plaintext3, 8);
	sg_set_buf(sg_src + 3, plaintext4, 16);
	sg_set_buf(sg_src + 4, plaintext5, 8);

	ahash_request_set_crypt(req, sg_src, result, len);
	crypto_ahash_digest(req);

	init_completion(&com);
	wait_for_completion(&com);

	test_dbg("sha256 test done.\n");

	ahash_request_free(req);
	crypto_free_ahash(tfm);
	return 0;
}

static int __init crypto_test_init(void)
{
	crypto_sha_test();
	crypto_aes_test();
	return 0;
}

static void __exit crypto_test_exit(void) 
{
}

module_init(crypto_test_init);
module_exit(crypto_test_exit);

MODULE_AUTHOR("Jian Tu <tujian2bgs@gmail.com>");
MODULE_DESCRIPTION("Crypto SHA,AES test module");
MODULE_LICENSE("GPL");

